import mmcv
import torch
import cv2
import numpy as np
from mmcv.runner import load_checkpoint
from Dataset import build_dataset
from mmcv.parallel import MMDataParallel
from mmdet.core import wrap_fp16_model
from mmdet.datasets import build_dataloader
from mmdet.models import build_detector
from mmdet.core import tensor2imgs


def Get_result(module, data, result, dataset=None, score_thr=0.3):
    """
    输入图像及筛查结果boxes, 画图并输出结果图;
    :param module: 乳腺筛查模型
    :param data: 输入筛查数据
    :param result: 筛查结果boxes
    :return: 输出筛查结果图
    """
    if isinstance(result, tuple):
        bbox_result, segm_result = result
    else:
        bbox_result, segm_result = result, None

    img_tensor = data['img'][0]
    img_metas = data['img_meta'][0].data[0]
    imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
    assert len(imgs) == len(img_metas)

    class_names = module.CLASSES

    for img, img_meta in zip(imgs, img_metas):
        h, w, _ = img_meta['img_shape']
        img_show = img[:h, :w, :]

        bboxes = np.vstack(bbox_result)
        # draw bounding boxes
        labels = [
            np.full(bbox.shape[0], i, dtype=np.int32)
            for i, bbox in enumerate(bbox_result)
        ]
        labels = np.concatenate(labels)
        # 此处修改mmcv.imshow_det_bboxes内puttext参数
        # cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 10),
        #                     cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color,3)
        img_show = mmcv.imshow_det_bboxes(
            img_show,
            bboxes,
            labels,
            class_names=class_names,
            score_thr=score_thr,
            thickness=5,
            font_scale=1.25,
            show=False)
    return img_show.astype(np.uint8)


def single_gpu_test(model, data_loader, show=False):
    """
    输入模型及图像列表,返回筛查结果图
    :param model: 筛查模型
    :param data_loader: 图像列表
    :return: 筛查结果图
    """
    model.eval()
    results = []
    dataset = data_loader.dataset
    prog_bar = mmcv.ProgressBar(len(dataset))
    for i, data in enumerate(data_loader):
        with torch.no_grad():
            result = model(return_loss=False, rescale=show, **data)
        if show:
            model.module.show_result(data, result)
        img = Get_result(model.module, data, result)
        results.append(img)
        batch_size = data['img'][0].size(0)
        for _ in range(batch_size):
            prog_bar.update()
    return results


class mmDetector(object):
    """
    init()初始化模型
    infrence()构建dataset,并对dataset进行筛查
        Return:筛查结果图
    """
    def __init__(self, config_file, show=False):
        self.cfg = mmcv.Config.fromfile(config_file)
        if self.cfg.get('cudnn_benchmark', False):
            torch.backends.cudnn.benchmark = True
        self.cfg.model.pretrained = None
        self.cfg.data.test.test_mode = True
        self.show = False

    def init(self):
        model = build_detector(self.cfg.model, train_cfg=None, test_cfg=self.cfg.test_cfg)
        fp16_cfg = self.cfg.get('fp16', None)
        if fp16_cfg is not None:
            wrap_fp16_model(model)
        self.model = MMDataParallel(model, device_ids=[0])
        checkpoint = load_checkpoint(model, str(self.cfg.model_root), map_location='cpu')
        model.CLASSES = self.cfg.CLASSES

    def infrence(self, mg_dcms):
        dataset = build_dataset(self.cfg.data.test, mg_dcms)
        data_loader = build_dataloader(
            dataset,
            imgs_per_gpu=1,
            workers_per_gpu=self.cfg.data.workers_per_gpu,
            shuffle=False)
        outputs = single_gpu_test(self.model, data_loader, self.show)
        return outputs

